import time, os
import torch
import numpy as np
from importlib import import_module
from parse import args

if __name__ == '__main__':
    if not os.path.exists(f"{args.dataset}"): raise FileNotFoundError("Data folder must exist!")
    if not os.path.exists(f"{args.dataset}/saved_dict"): os.mkdir(f"{args.dataset}/saved_dict")

    # 搜狗新闻:embedding_SougouNews.npz, 腾讯:embedding_Tencent.npz, 随机初始化:random
    if args.model == 'FastText':
        import utils_fasttext as utils
        args.embedding = 'random'
    elif 'bert' in args.model:
        import utils_bert as utils
    else:
        import utils as utils

    if args.multi_labels:
        from train_eval import train_multi_labels as train, init_network, infer_multi_labels as infer
    else:
        from train_eval import train, init_network, infer

    import_model = import_module('models.' + args.model)
    config = import_model.Config(args)
    if args.MLB_path: config.MLB_path = args.dataset + "/data/" + args.MLB_path

    start_time = time.time()
    print("Loading data...")
    if args.infer_mode:
        vocab, infer_data = utils.build_dataset(config, args)
        infer_iter = utils.build_iterator(infer_data, config)
        config.n_vocab = len(vocab)
        model = import_model.Model(config).to(config.device)
        model.load_state_dict(torch.load(config.save_path))
        time_dif = utils.get_time_dif(start_time)
        print("Time usage:", time_dif)
        infer(config, model, infer_iter, args)

    else:
        vocab, train_data, dev_data, test_data = utils.build_dataset(config, args)
        train_iter = utils.build_iterator(train_data, config)
        dev_iter = utils.build_iterator(dev_data, config)
        test_iter = utils.build_iterator(test_data, config)
        time_dif = utils.get_time_dif(start_time)
        print("Time usage:", time_dif)

        # train
        config.n_vocab = len(vocab)
        model = import_model.Model(config).to(config.device)
        if args.model != 'Transformer':
            init_network(model, args=args)
        # print(model.parameters)
        train(config, model, train_iter, dev_iter, test_iter, args)

